import time
import uuid
import argparse
import datetime
import numpy as np
from tqdm import tqdm
import os
import torch
import torch.nn.functional as F
from torch_sparse import remove_diag

from model import *
from utils import *
from arch import archs
from dataProcess import *

def main(args):
    if args.seed > 0:
        set_random_seed(args.seed)

    # 加载数据集，返回图g、邻接矩阵adjs、初始标签init_labels、类别数量num_classes、数据加载器dl、训练验证节点trainval_nid和测试节点test_nid。
    # todo 暂时不太清楚用途：test_nid_full  多余的节点
    g, adjs, init_labels, num_classes, dl, train_nid, val_nid, test_nid, test_nid_full = load_dataset(args)

    # 这句代码的整体作用是将邻接矩阵中每条边的权重重新设置为该边起始节点度的倒数。这样做的目的是进行归一化处理，使得每个节点连接的边权重之和为1，通常用于消息传递中的归一化步骤。
    adjs = adjsProcess(adjs)

    # =======
    # rearange node idx (for feats & labels)
    # =======
    train_node_nums = len(train_nid)
    valid_node_nums = len(val_nid)
    test_node_nums = len(test_nid)

    # 设置训练和验证分割点
    trainval_point = train_node_nums
    # 设置验证和测试分割点
    valtest_point = trainval_point + valid_node_nums
    # 计算训练、验证和测试节点的总数
    total_num_nodes = train_node_nums + valid_node_nums + test_node_nums

    #目标节点数量
    num_nodes = dl.nodes['count'][0]

    # todo 不确定哪个数据集会有用
    # 检查 total_num_nodes 是否小于 num_nodes，如果用于学习的总节点数量小于目标节点的数量，则表示有多余节点。
    # # 检查是否有额外的未标注节点，如果存在，则记录这些节点。
    if total_num_nodes < num_nodes:
        # 创建一个布尔数组 flag，初始化为全 True，表示假设所有节点都是多余的。
        flag = np.ones(num_nodes, dtype=bool)
        # 将训练节点的索引对应位置设为 False，表示这些节点不是多余的。
        flag[train_nid] = 0
        # 将验证节点索引对应位置设为 False。
        flag[val_nid] = 0
        # 将测试节点索引对应位置设为 False。
        flag[test_nid] = 0
        # 获取 flag 数组中值为 True 的位置索引，这些即为多余节点的索引。
        extra_nid = np.where(flag)[0]
        # 打印出找到的多余节点数量 len(extra_nid)。
        print(f'Find {len(extra_nid)} extra nid for dataset {args.dataset}')
    else:
        # 如果 total_num_nodes 大于等于 num_nodes，则没有多余节点，将 extra_nid 设置为空数组
        extra_nid = np.array([])

    # init2sort 是一个 torch.LongTensor 类型的张量，将 train_nid、val_nid、test_nid 和 extra_nid 连接在一起，生成了包含所有节点的一个序列，顺序是训练、验证、测试和多余节点。
    init2sort = torch.LongTensor(np.concatenate([train_nid, val_nid, test_nid, extra_nid]))
    # sort2init 表示原始索引到排序索引的映射关系。
    # torch.argsort(init2sort) 返回 init2sort 排序后对应的原始索引位置
    sort2init = torch.argsort(init2sort)
    # 断言语句用于验证重新排列后的标签与初始标签 init_labels 相同。
    # init_labels[init2sort][sort2init] 先按 init2sort 的顺序提取标签，然后按 sort2init 的顺序重新排列，还原成 init_labels，以确保顺序恢复后数据没有发生改变。
    assert torch.all(init_labels[init2sort][sort2init] == init_labels)
    # labels 是按 init2sort 排序后的标签，用于模型训练中节点标签的有序输入。
    labels = init_labels[init2sort]

    # =======
    # neighbor aggregation
    # =======
    # 根据数据集的不同，设置目标节点类型tgt_type、所有节点类型node_types和额外元路径extra_metapath。
    if args.dataset == 'DBLP':
        tgt_type = 'A'
        node_types = ['A', 'P', 'T', 'V']
        extra_metapath = []
    elif args.dataset == 'ACM':
        tgt_type = 'P'
        node_types = ['P', 'A', 'C']
        extra_metapath = []
    elif args.dataset == 'IMDB':
        tgt_type = 'M'
        node_types = ['M', 'A', 'D', 'K']
        extra_metapath = []
    elif args.dataset == 'Freebase':
        tgt_type = '0'
        node_types = [str(i) for i in range(8)]
        extra_metapath = []
    else:
        assert 0
    # 过滤extra_metapath，只保留长度大于args.num_hops + 1的路径。
    extra_metapath = [ele for ele in extra_metapath if len(ele) > args.num_hops + 1]

    # 打印当前特征传播的跳数args.num_hops。
    print(f'Current num hops = {args.num_hops}')

    # 特征传播设备为cpu
    prop_device = 'cpu'
    # 设置存储设备为CPU。
    store_device = 'cpu'

    # compute k-hop feature
    feats, data_size = neighboorAggregation(args, g, dl, extra_metapath, tgt_type, init2sort)

    # =======
    # 设置检查点文件夹
    checkpt_folder = f'./output/{args.dataset}/'
    if not os.path.exists(checkpt_folder):
        os.makedirs(checkpt_folder)
    checkpt_file = checkpt_folder + uuid.uuid4().hex

    # 如果 args.amp 为真，则初始化 GradScaler，用于混合精度训练。
    if args.amp:
        scalar = torch.cuda.amp.GradScaler()
    else:
        scalar = None

    # 选择设备
    device = 'cuda:{}'.format(args.gpu) if not args.cpu else 'cpu'
    # 对于非 IMDB 数据集，将标签转换为 long 类型，并移动到选定设备。
    # 对于 IMDB 数据集，将标签转换为 float 类型，并移动到选定设备。
    if args.dataset != 'IMDB':
        labels_cuda = labels.long().to(device)
    else:
        labels = labels.float()
        labels_cuda = labels.to(device)

    # 开始训练阶段（这里只有 stage=0）。
    for stage in [0]:
        # epochs 表示训练的轮数，取决于 args.stage 参数。
        epochs = args.stage

        # =======
        # labels propagate alongside the metapath
        # =======
        label_feats = {}
        # 根据 args.label_feats 参数，决定是否进行标签传播
        if args.label_feats:
            # 如果数据集不是 IMDB，对训练集标签进行 one-hot 编码，构建 label_onehot。
            # 如果数据集是 IMDB，则直接使用初始化的标签，赋值给 label_onehot。
            if args.dataset != 'IMDB':
                label_onehot = torch.zeros((num_nodes, num_classes))
                # init_labels[train_nid] 作为输入，每个标签值（整数）会被转换为长度为 num_classes 的独热向量。例如，如果 num_classes 为 4，且 init_labels[train_nid] 中的某个标签为 2，则该标签将被转换为 [0, 0, 1, 0]。
                # 结果是一个大小为 (len(train_nid), num_classes) 的张量，表示训练集中所有节点的独热编码标签。
                label_onehot[train_nid] = F.one_hot(init_labels[train_nid], num_classes).float()
            else:
                label_onehot = torch.zeros((num_nodes, num_classes))
                label_onehot[train_nid] = init_labels[train_nid].float()

            # 根据不同数据集设置 extra_metapath（额外元路径）。在这里均为空列表。
            # 如果数据集不在上述列表中，则终止程序并抛出异常。
            if args.dataset == 'DBLP':
                extra_metapath = []
            elif args.dataset == 'IMDB':
                extra_metapath = []
            elif args.dataset == 'ACM':
                extra_metapath = []
            else:
                assert 0

            # 过滤 extra_metapath，保留路径长度大于 args.num_label_hops + 1 的路径。
            # max_length 是标签传播的最大路径长度，取 args.num_label_hops + 1 和 extra_metapath 中最长路径的最大值。
            extra_metapath = [ele for ele in extra_metapath if len(ele) > args.num_label_hops + 1]
            if len(extra_metapath):
                max_length = max(args.num_label_hops + 1, max([len(ele) for ele in extra_metapath]))
            else:
                max_length = args.num_label_hops + 1
            # 打印标签传播的跳数，用于调试或检查。
            print(f'Current label-prop num hops = {args.num_label_hops}')

            # compute k-hop feature
            prop_tic = datetime.datetime.now()
            # 调用 hg_propagate_sparse_pyg 函数，基于指定的 args.num_label_hops 和 max_length 进行稀疏邻接矩阵传播，获取 meta_adjs。
            meta_adjs = hg_propagate_sparse_pyg(
                adjs, tgt_type, args.num_label_hops, max_length, extra_metapath, prop_feats=False, echo=False, prop_device=prop_device)
            # 打印每个元路径邻接矩阵的键名和尺寸信息。
            print(f'For label propagation, meta_adjs: (in SparseTensor mode)')
            for k, v in meta_adjs.items():
                print(k, v.sizes())
            print()

            # 遍历 meta_adjs，对每个邻接矩阵 v 进行处理，移除对角线元素 remove_diag(v)。
            # 乘以 label_onehot，将标签特征存入 label_feats。
            for k, v in tqdm(meta_adjs.items()):
                label_feats[k] = remove_diag(v) @ label_onehot
            # 手动触发垃圾回收，释放不必要的内存。
            gc.collect()


            # 检查标签特征的准确性，调用 check_acc 函数评估标签传播的结果。
            # 若数据集为 IMDB，设定损失类型为二元交叉熵 bce，否则采用默认损失。
            if args.dataset == 'IMDB':
                condition = lambda ra,rb,rc,k: True
                check_acc(label_feats, condition, init_labels, train_nid, val_nid, test_nid, show_test=False, loss_type='bce')
            else:
                condition = lambda ra,rb,rc,k: True
                check_acc(label_feats, condition, init_labels, train_nid, val_nid, test_nid, show_test=True)
            print('Involved label keys', label_feats.keys())

            # 将 label_feats 中的每个特征根据 init2sort 重新排序，并保留符合 archs[args.arch][1] 中的特征键。
            label_feats = {k: v[init2sort] for k,v in label_feats.items() if k in archs[args.arch][1]}   # if k in archs[args.arch][1]

            prop_toc = datetime.datetime.now()
            print(f'Time used for label prop {prop_toc - prop_tic}')

        # =======
        # Train & eval loaders
        # =======
        # 创建训练数据加载器 train_loader，随机打乱训练节点索引，并设定 batch_size 和 drop_last 参数。
        train_loader = torch.utils.data.DataLoader(
            torch.arange(train_node_nums), batch_size=args.batch_size, shuffle=True, drop_last=False)

        # =======
        # Mask & Smooth
        # =======
        with_mask = False

        # 初始化评估数据加载器 eval_loader 和多余数据加载器 full_loader。
        eval_loader, full_loader = [], []
        # 设置 batchsize 为 args.batch_size 的两倍。
        # todo 这里不明白为什么设置为2倍
        batchsize = 2 * args.batch_size

        # 将训练、验证和测试节点分批加载到 eval_loader。
        # 对每个批次索引 batch_idx，计算起止点，将相应特征和标签特征提取并添加到 eval_loader。
        for batch_idx in range((total_num_nodes-1) // batchsize + 1):
            batch_start = batch_idx * batchsize
            batch_end = min(total_num_nodes, (batch_idx+1) * batchsize)
            batch = torch.arange(batch_start, batch_end)

            batch_feats = {k: x[batch_start:batch_end] for k, x in feats.items()}
            batch_labels_feats = {k: x[batch_start:batch_end] for k, x in label_feats.items()}

            batch_mask = None
            eval_loader.append((batch, batch_feats, batch_labels_feats, batch_mask))

        # 多余数据（额外节点）
        for batch_idx in range((num_nodes-total_num_nodes-1) // batchsize + 1):
            batch_start = batch_idx * batchsize + total_num_nodes
            batch_end = min(num_nodes, (batch_idx+1) * batchsize + total_num_nodes)
            batch = torch.arange(batch_start, batch_end)

            batch_feats = {k: x[batch_start:batch_end] for k, x in feats.items()}
            batch_labels_feats = {k: x[batch_start:batch_end] for k, x in label_feats.items()}

            batch_mask = None
            full_loader.append((batch, batch_feats, batch_labels_feats, batch_mask))

        # =======
        # Construct network
        # =======
        # 释放显存和手动触发垃圾回收。
        torch.cuda.empty_cache()
        gc.collect()

        model = LMSPS(args.embed_size, args.hidden, num_classes, feats.keys(), label_feats.keys(), tgt_type,
            args.dropout, args.input_drop, args.att_drop, args.label_drop,
            args.n_layers_2,  args.residual, bns=args.bns, data_size=data_size, path=archs[args.arch][0],
            label_path=archs[args.arch][1], eps=args.eps, device=device)
        
        model = model.to(device)

        # 打印模型参数总数（在 args.seed 为初始值时）。
        if args.seed == args.seeds[0]:
            #print(model)
            print("# Params:", get_n_params(model))

        # # 对于 IMDB 数据集，使用二元交叉熵损失 BCEWithLogitsLoss。对于其他数据集，使用交叉熵损失 CrossEntropyLoss。
        if args.dataset == 'IMDB':
            loss_fcn = nn.BCEWithLogitsLoss()
        else:
            loss_fcn = nn.CrossEntropyLoss()
        # 创建 Adam 优化器，使用模型的参数，并设置学习率和权重衰减。
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
                                    weight_decay=args.weight_decay)

        # 初始化一些变量以跟踪最佳epoch、最佳验证损失、最佳测试损失、验证和测试精度列表等。
        best_epoch = -1
        best_val_loss = 1000000
        best_test_loss = 0
        best_val = (0,0)
        best_test = (0,0)
        val_loss_list, test_loss_list = [], []
        val_acc_list, test_acc_list = [], []
        actual_loss_list, actual_acc_list = [], []
        store_list = []
        best_pred = None
        count = 0

        train_times = []

        # 开始训练循环，迭代 args.stage 次。tqdm 用于显示进度条。
        for epoch in tqdm(range(args.stage)):
            # 手动调用垃圾回收器，确保没有未清理的内存。然后同步 CUDA 设备，记录当前时间。
            gc.collect()
            if not args.cpu: torch.cuda.synchronize()
            start = time.time()

            # 调用 train 函数进行一次训练，返回损失和准确率。
            loss, acc = train(model, feats, label_feats, labels_cuda, loss_fcn, optimizer, train_loader, evaluator, scalar=scalar)

            # 再次同步 CUDA 设备，记录训练结束的时间。
            if not args.cpu: torch.cuda.synchronize()
            end = time.time()

            log = "Epoch {}, training Time(s): {:.4f}, estimated train loss {:.4f}, acc {:.4f}, {:.4f}\n".format(epoch, end - start,loss, acc[0]*100, acc[1]*100)
            if not args.cpu: torch.cuda.empty_cache()
            # 将本轮训练所需的时间添加到 train_times 列表中。
            train_times.append(end-start)

            start = time.time()
            # 在验证集和测试集上评估模型性能。
            # 记录验证损失和测试损失，并基于它们更新最佳模型。
            # 禁用梯度计算以提高评估速度，将模型设置为评估模式，并初始化 raw_preds 列表用于存储预测结果。
            with torch.no_grad():
                model.eval()
                raw_preds = []

                # 对于每个评估批次，将数据转移到设备并执行模型的前向传播，生成预测结果并将其添加到 raw_preds 列表。
                for batch, batch_feats, batch_labels_feats, batch_mask in eval_loader:
                    batch = batch.to(device)
                    batch_feats = {k: x.to(device) for k, x in batch_feats.items()}
                    batch_labels_feats = {k: x.to(device) for k, x in batch_labels_feats.items()}
                    if with_mask:
                        batch_mask = {k: x.to(device) for k, x in batch_mask.items()}
                    else:
                        batch_mask = None
                    raw_preds.append(model(batch, batch_feats, batch_labels_feats, batch_mask).cpu())

                # 将所有批次的预测结果合并为一个张量。
                raw_preds = torch.cat(raw_preds, dim=0)
                # 计算训练、验证和测试损失。
                loss_train = loss_fcn(raw_preds[:trainval_point], labels[:trainval_point]).item()
                loss_val = loss_fcn(raw_preds[trainval_point:valtest_point], labels[trainval_point:valtest_point]).item()
                loss_test = loss_fcn(raw_preds[valtest_point:total_num_nodes], labels[valtest_point:total_num_nodes]).item()

            # 根据数据集类型，生成最终预测结果。
            if args.dataset != 'IMDB':
                preds = raw_preds.argmax(dim=-1)
            else:
                preds = (raw_preds > 0.).int()

            # 使用 evaluator 计算训练、验证和测试集的准确率。
            train_acc = evaluator(preds[:trainval_point], labels[:trainval_point])
            val_acc = evaluator(preds[trainval_point:valtest_point], labels[trainval_point:valtest_point])
            test_acc = evaluator(preds[valtest_point:total_num_nodes], labels[valtest_point:total_num_nodes])

            end = time.time()

            #log += f'evaluation Time: {end-start}, Train loss: {loss_train}, Val loss: {loss_val}, Test loss: {loss_test}\n'
            log += f'Val loss: {loss_val}, Test loss: {loss_test}\n'
            log += 'Train acc: ({:.4f}, {:.4f}), Val acc: ({:.4f}, {:.4f}), Test acc: ({:.4f}, {:.4f}) ({})\n'.format(
                train_acc[0]*100, train_acc[1]*100, val_acc[0]*100, val_acc[1]*100, test_acc[0]*100, test_acc[1]*100, total_num_nodes-valtest_point)

            # 如果当前验证损失是最好的，则更新最佳模型参数。如果连续多轮没有改进，则可能会提前停止训练。
            if loss_val <= best_val_loss:
                best_epoch = epoch
                best_val_loss = loss_val
                best_test_loss = loss_test
                best_val = val_acc
                best_test = test_acc

                best_pred = raw_preds
                torch.save(model.state_dict(), f'{checkpt_file}.pkl')
                wait = 0
            else:
                wait += 1
                if wait == args.patience:
                    break

            # 每隔 10 轮输出当前最佳验证损失和测试损失。
            if epoch > 0 and epoch % 10 == 0: 
                 log += f'\tCurrent best at epoch {best_epoch} with Val loss {best_val_loss:.4f} ({best_val[0]*100:.4f}, {best_val[1]*100:.4f})' \
                    + f', Test loss {best_test_loss:.4f} ({best_test[0]*100:.4f}, {best_test[1]*100:.4f})'
            print(log)

        print('average train times', sum(train_times) / len(train_times))

        print(f'Best Epoch {best_epoch} at {checkpt_file.split("/")[-1]}\n\tFinal Val loss {best_val_loss:.4f} ({best_val[0]*100:.4f}, {best_val[1]*100:.4f})'
            + f', Test loss {best_test_loss:.4f} ({best_test[0]*100:.4f}, {best_test[1]*100:.4f})')   # micro  macro


        # 如果有 full_loader，则加载保存的最佳模型，进行评估并生成最终预测。
        if len(full_loader):
            model.load_state_dict(torch.load(f'{checkpt_file}.pkl', map_location='cpu'), strict=True)
            torch.cuda.empty_cache()
            with torch.no_grad():
                model.eval()
                raw_preds = []

                for batch, batch_feats, batch_labels_feats, batch_mask in full_loader:
                    batch = batch.to(device)
                    batch_feats = {k: x.to(device) for k, x in batch_feats.items()}
                    batch_labels_feats = {k: x.to(device) for k, x in batch_labels_feats.items()}
                    if with_mask:
                        batch_mask = {k: x.to(device) for k, x in batch_mask.items()}
                    else:
                        batch_mask = None
                    raw_preds.append(model(batch, batch_feats, batch_labels_feats, batch_mask).cpu())
                raw_preds = torch.cat(raw_preds, dim=0)
            best_pred = torch.cat((best_pred, raw_preds), dim=0)

        #torch.save(best_pred, f'{checkpt_file}.pt')

        # 根据数据集类型生成最终预测概率并生成输出文件，保存预测结果以供后续评估。
        if args.dataset != 'IMDB':
            predict_prob = best_pred.softmax(dim=1)
        else:
            predict_prob = torch.sigmoid(best_pred)

        test_logits = predict_prob[sort2init][test_nid_full]
        if args.dataset != 'IMDB':
            pred = test_logits.cpu().numpy().argmax(axis=1)
            dl.gen_file_for_evaluate(test_idx=test_nid_full, label=pred, file_path=f"./output/{args.dataset}_{args.seed}_{checkpt_file.split('/')[-1]}.txt")
        else:
            pred = (test_logits.cpu().numpy()>0.5).astype(int)
            dl.gen_file_for_evaluate(test_idx=test_nid_full, label=pred, file_path=f"./output/{args.dataset}_{args.seed}_{checkpt_file.split('/')[-1]}.txt", mode='multi')

    # # 计算并打印训练、验证和测试集的准确率，以便评估模型在各个数据集上的表现。
    if args.dataset != 'IMDB':
        preds = predict_prob.argmax(dim=1, keepdim=True)
    else:
        preds = (predict_prob > 0.5).int()
    train_acc = evaluator(labels[:trainval_point], preds[:trainval_point])
    val_acc = evaluator(labels[trainval_point:valtest_point], preds[trainval_point:valtest_point])
    test_acc = evaluator(labels[valtest_point:total_num_nodes], preds[valtest_point:total_num_nodes])

    print(f'train_acc ({train_acc[0]*100:.2f}, {train_acc[1]*100:.2f}) ' \
        + f'val_acc ({val_acc[0]*100:.2f}, {val_acc[1]*100:.2f}) ' \
        + f'test_acc ({test_acc[0]*100:.2f}, {test_acc[1]*100:.2f})')
    print(checkpt_file.split('/')[-1])

    # 删除保存的最优的模型
    os.remove(f'{checkpt_file}.pkl')

    return [test_acc[0]*100, test_acc[1]*100]

def parse_args(args=None):
    parser = argparse.ArgumentParser(description='LMSPS')
    ## For environment costruction
    parser.add_argument("--seeds", nargs='+', type=int, default=[1],
                        help="the seed used in the training")
    parser.add_argument("--dataset", type=str, default="ogbn-mag")
    parser.add_argument("--gpu", type=int, default=0)
    parser.add_argument("--cpu", action='store_true', default=False)
    parser.add_argument("--root", type=str, default="../data/")
    parser.add_argument("--stage", type=int, default=200, help="The epoch setting for each stage.")
    parser.add_argument("--embed-size", type=int, default=512,
                        help="inital embedding size of nodes with no attributes")
    parser.add_argument("--num-hops", type=int, default=2,
                        help="number of hops for propagation of raw labels")
    parser.add_argument("--label-feats", action='store_true', default=False,
                        help="whether to use the label propagated features")
    parser.add_argument("--num-label-hops", type=int, default=2,
                        help="number of hops for propagation of raw features")
    ## For network structure
    parser.add_argument("--hidden", type=int, default=512)
    parser.add_argument("--dropout", type=float, default=0.45,
                        help="dropout on activation")
    parser.add_argument("--n-layers-2", type=int, default=3,
                        help="number of layers of the downstream task")
    parser.add_argument("--input-drop", type=float, default=0.1,
                        help="input dropout of input features")
    parser.add_argument("--att-drop", type=float, default=0.,
                        help="attention dropout of model")
    parser.add_argument("--label-drop", type=float, default=0.,
                        help="label feature dropout of model")
    parser.add_argument("--residual", action='store_true', default=False,
                        help="whether to add residual branch the raw input features")
    parser.add_argument("--bns", action='store_true', default=False,
                        help="whether to process the input features")
    ## for training
    parser.add_argument("--amp", action='store_true', default=False,
                        help="whether to amp to accelerate training with float16(half) calculation")
    parser.add_argument("--lr", type=float, default=0.001)
    parser.add_argument("--weight-decay", type=float, default=0)
    parser.add_argument("--eval-every", type=int, default=1)
    parser.add_argument("--batch-size", type=int, default=10000)
    parser.add_argument("--patience", type=int, default=50,
                        help="early stop of times of the experiment")
    parser.add_argument("--edge_mask_ratio", type=float, default=0)
    parser.add_argument('--arch', type=str, default='DBLP')
    parser.add_argument("--eps", type=float, default=0)   #1e-12
    parser.add_argument("--ACM_keep_F", action='store_true', default=False,
                        help="whether to use Field type")

    return parser.parse_args(args)

if __name__ == '__main__':
    args = parse_args()

    # if args.dataset == 'ACM':
    #     args.ACM_keep_F = False

    args.seed = args.seeds[0]
    print(args)

    results = []
    for seed in args.seeds:
        args.seed = seed
        print('Restart with seed =', seed)
        result = main(args)
        results.append(result)
    print('results', results)
    

    if args.dataset == 'IMDB':
        results.sort(key=lambda x: x[0], reverse=True)
        print(results)
        results = results[:5]
        mima = list(map(list, zip(*results)))
        #print(f'micro: {mima[0]}', f'macro: {mima[1]}')
        print(f'micro_mean: {np.mean(mima[0]):.2f}', f'micro_std: {np.std(mima[0]):.2f}')
        print(f'macro_mean: {np.mean(mima[1]):.2f}', f'macro_std: {np.std(mima[1]):.2f}')
    else:
        aver = list(map(list, zip(*results)))
        print(f'micro_aver: {np.mean(aver[0]):.2f}', f'micro_std: {np.std(aver[0]):.2f}')
        print(f'macro_aver: {np.mean(aver[1]):.2f}', f'macro_std: {np.std(aver[1]):.2f}')

